package com.mashibing.loadbalancer;

import org.springframework.beans.factory.ObjectProvider;
import org.springframework.cloud.client.ServiceInstance;
import org.springframework.cloud.client.loadbalancer.DefaultResponse;
import org.springframework.cloud.client.loadbalancer.Request;
import org.springframework.cloud.client.loadbalancer.RequestDataContext;
import org.springframework.cloud.client.loadbalancer.Response;
import org.springframework.cloud.loadbalancer.core.NoopServiceInstanceListSupplier;
import org.springframework.cloud.loadbalancer.core.ReactorServiceInstanceLoadBalancer;
import org.springframework.cloud.loadbalancer.core.ServiceInstanceListSupplier;
import reactor.core.publisher.Mono;

import java.util.List;
import java.util.Random;
import java.util.concurrent.atomic.AtomicInteger;

public class TagLoadBalancer implements ReactorServiceInstanceLoadBalancer {

    String serviceId;

    ObjectProvider<ServiceInstanceListSupplier> serviceInstanceListProviders;

    final String CUSTOMIZED_TAG = "tag";
    final String VERSION = "version";
    final String V1 = "v1";
    final String V2 = "v2";
    final String NUM = "num";
    final String STR = "str";

    public TagLoadBalancer(String serviceId, ObjectProvider<ServiceInstanceListSupplier> serviceInstanceListProviders) {
        this.serviceId = serviceId;
        this.serviceInstanceListProviders = serviceInstanceListProviders;
    }

    @Override
    public Mono<Response<ServiceInstance>> choose(Request request) {
        ServiceInstanceListSupplier ifAvailable = serviceInstanceListProviders.getIfAvailable(NoopServiceInstanceListSupplier::new);
        Mono<List<ServiceInstance>> next = ifAvailable.get(request).next();
        return next.mapNotNull(serviceInstances -> this.getInstanceResponse(request, serviceInstances));
    }

    private Response<ServiceInstance> getInstanceResponse(Request<?> request, List<ServiceInstance> serviceInstances) {

        ServiceInstance resultServiceInstance = null;
        List<String> tags = ((RequestDataContext) (request.getContext())).getClientRequest().getHeaders().get(CUSTOMIZED_TAG);
        if (tags == null) {
            int seedPosition = new Random().nextInt(1000);
            AtomicInteger position = new AtomicInteger(seedPosition);

            int pos = Math.abs(position.incrementAndGet());

            resultServiceInstance = serviceInstances.get(pos % serviceInstances.size());
        } else {
            String tag = tags.get(0);
            for (ServiceInstance serviceInstance : serviceInstances) {
                String version = serviceInstance.getMetadata().get(VERSION);
                if ((tag.equals(NUM) && version.equals(V1)) || (tag.equals(STR) && version.equals(V2))) {
                    resultServiceInstance = serviceInstance;
                    break;
                }
            }
        }
        return new DefaultResponse((resultServiceInstance));
    }
}
